
#pragma once

#include <algorithm>
#include <stack>
#include <list>
#include "Scene.h"
#include "SceneRenderer.h"
#include "Surfaces.h"
#include "Camera.h"
#include "Ray.h"
#include "Material.h"
#include "Profiler.h"
#include "BVH.h"
#include "Prerequisites.h"

class BVHRenderer : public SceneRenderer
{
protected:
	BVHNode<DefaultBV, Surface> *mRootNode;
	mutable int mRayCounter;
	mutable int mShadowRayCounter;
	bool mSampleQuad;

	Ogre::ColourValue samplePixel(const Camera &camera, Scene &scene, int x, int y)
	{
		if (!mSampleQuad)
		{
			Ogre::ColourValue col;
			_raytrace(col, camera.getProjection().pixelToRay((float)x, (float)y), scene, 0);
			return col;
		}
		else
		{
			Ogre::ColourValue col0,col1, col2, col3;
			float offset = 0.25f;
			_raytrace(col0, camera.getProjection().pixelToRay(x+offset, y+offset), scene, 0);
			_raytrace(col1, camera.getProjection().pixelToRay(x+offset, y-offset), scene, 0);
			_raytrace(col2, camera.getProjection().pixelToRay(x-offset, y+offset), scene, 0);
			_raytrace(col3, camera.getProjection().pixelToRay(x-offset, y-offset), scene, 0);
			return (col0+col1+col2+col3) * 0.25f;
		}
	}

	/// Non-virtual method performance for a few miliseconds
	__forceinline bool _raytrace(Ogre::ColourValue &outColor, const Ray &ray, const Scene &scene, int recursionLevel) const
	{		
		mRayCounter++;
		Ray::Intersection intersection;
		const Surface *surface = mRootNode->rayIntersectClosest(intersection, ray);
		if (surface)
		{
			FragmentProgramData fragData;
			fragData.position = ray.origin + intersection.t*ray.direction;
			surface->computeFragmentProgramData(fragData, intersection);
			(surface->getMaterial()->fn)(outColor, fragData, *surface->getMaterial(), scene, *this, recursionLevel);
		}
		else outColor = scene.backgroundColor;

		return (surface != nullptr);
	}

public:
	virtual ~BVHRenderer()
	{
		if (mRootNode) delete mRootNode;
	}
	BVHRenderer()
	{
		mRootNode = nullptr;
		mRayCounter = 0;
		mShadowRayCounter = 0;
		mSampleQuad = false;
	}

	void setQuadSampling(bool sampleQuad)
	{
		mSampleQuad = sampleQuad;
	}

	bool raytrace(Ogre::ColourValue &outColor, const Ray &ray, const Scene &scene, int recursionLevel) const
	{
		return _raytrace(outColor, ray, scene, recursionLevel);
	}

	virtual bool raycastAny(const Ray &ray, const Scene &scene, float maxDist) const
	{
		mShadowRayCounter++;
		return mRootNode->rayIntersectAny(ray, maxDist) != nullptr;
	}

	virtual void renderScene(std::vector<Ogre::ColourValue> &renderedImage, Scene &scene, const Camera &camera)
	{
		mRayCounter = 0;
		mShadowRayCounter = 0;
		renderedImage.clear();
		renderedImage.resize(camera.getProjection().getWidth() * camera.getProjection().getHeight());

		Profiler::Timestamp timeStamp = Profiler::timestamp();
		scene.applyViewTransform(camera.getViewMatrix());
		Profiler::printJobDuration("Scene preprocessing", timeStamp);

		//construct the BVH
		timeStamp = Profiler::timestamp();

		if (mRootNode) delete mRootNode;
		std::vector<Surface*> surfaces;

		for (auto i = scene.spheres.begin(); i != scene.spheres.end(); i++) surfaces.push_back(*i);

		for (auto iMesh = scene.meshes.begin(); iMesh != scene.meshes.end(); ++iMesh)
		{
			for (auto iTri = (*iMesh)->triangleSurfaces.begin(); iTri != (*iMesh)->triangleSurfaces.end(); iTri++)
				surfaces.push_back(&(*iTri));
		}
		mRootNode = new BVHNode<DefaultBV, Surface>(surfaces, 0, surfaces.size(), 0);
		Profiler::printJobDuration("BVH creation", timeStamp);
		std::cout << "Tree height max: " << mRootNode->getHeight() << std::endl;
		std::cout << "Tree height avg: " << mRootNode->getHeightAvg() << std::endl;

		timeStamp = Profiler::timestamp();
		int pixelIndex = 0;
		for (int y = camera.getProjection().getHeight()-1; y >= 0 ; y--)
		{
			for (int x = 0; x < camera.getProjection().getWidth(); x++)
			{
				renderedImage[pixelIndex] = samplePixel(camera, scene, x, y);
				renderedImage[pixelIndex].saturate();
				pixelIndex++;
			}
		}
		Profiler::printJobDuration("Ray tracing", timeStamp);
		std::cout << mRayCounter << " rays were casted." << std::endl;
		std::cout << mShadowRayCounter << " shadow rays were casted." << std::endl;
	}
};